<template>
  <div class="wh-full overflow-hidden flex-col gap-12px">
    <div class="flex justify-between">
      <n-input-group class="w-400px">
        <n-input v-model:value.trim="modelName" clearable />
        <n-button type="primary" @click="handlePull" :disabled="!modelName">下载模型</n-button>
      </n-input-group>
    </div>
    <n-data-table
      size="small"
      :columns="columns"
      :data="data"
      bordered
      striped
      :summary="summaryMethod"
      :single-line="false"
    />
  </div>
</template>
<script lang="ts" setup>
import { h, onMounted, ref } from 'vue'
import {
  NButton,
  NInput,
  NTag,
  useMessage,
  type DataTableColumn,
  type DataTableCreateSummary,
} from 'naive-ui'
import ollama from 'ollama'
import type { ModelResponse, ProgressResponse } from 'ollama'
import dayjs from 'dayjs'
import { useRouter } from 'vue-router'
import { useStorage } from '@vueuse/core'
import type { AbortableAsyncIterator } from 'ollama/src/utils.js'
defineOptions({ name: 'ModelList' })
const router = useRouter()
const message = useMessage()

interface ModelList extends ModelResponse {
  /**状态 */
  status: string
  /**状态信息 */
  msg: string
  /**已下载 */
  completed: number
  /**下载速度 */
  speed: number
  /**当前时间 */
  time: number
  /**取消器 */
  aborter?: AbortableAsyncIterator<ProgressResponse>
}

/**表头 */
const columns = ref<DataTableColumn<ModelList>[]>([
  { title: '名称', key: 'name', width: 350 },
  {
    title: '大小',
    key: 'size',
    render: (row) => (row.size / 1024 / 1024 / 1024).toFixed(2) + 'GB',
    width: 150,
  },
  {
    title: '更新时间',
    key: 'modified_at',
    render: ({ modified_at }) => dayjs(modified_at).format('YYYY-MM-DD HH:mm:ss'),
    width: 200,
  },
  {
    title: '状态',
    key: 'status',
    render: (row) => {
      if (row.status === 'success') {
        return h('div', { class: 'flex items-center' }, [
          h(NTag, { type: 'success' }, () => '已下载'),
          h(
            NButton,
            {
              text: true,
              type: 'primary',
              class: 'ml-8px',
              onClick: () => {
                pullFromOllama(row.name)
              },
            },
            () => '更新',
          ),
        ])
      }
      const speed =
        row.speed / 1024 > 1024
          ? (row.speed / 1024 / 1024).toFixed(2) + 'MB/s'
          : (row.speed / 1024).toFixed(2) + 'KB/s'
      return h(
        'div',
        { class: 'flex items-center' },
        h('div', { class: 'flex-c-c' }, [
          h(NButton, {
            text: true,
            class: 'text-gray px-4px',
            loading: row.status !== 'success' && row.status !== 'error',
          }),
          h('span', { class: 'px-4px color-primary' }, '当前文件' + row.status),
          h(
            'span',
            { class: 'px-4px color-primary' },
            row.completed / 1024 / 1024 > 1024
              ? '已完成：' + (row.completed / 1024 / 1024 / 1024).toFixed(2) + 'GB'
              : '已完成：' + (row.completed / 1024 / 1024).toFixed(2) + 'MB',
          ),
          h('span', { class: 'px-4px color-primary' }, '当前速度' + speed),
        ]),
      )
    },
  },
  {
    title: '操作',
    key: '',
    width: 100,
    render: (row) => {
      return h(
        NButton,
        {
          onClick: () => {
            if (row.status === 'success') {
              router.push({ path: `/chat` })
            } else {
              row.aborter?.abort()
            }
          },
          type: 'primary',
        },
        () => (row.status === 'success' ? '聊天' : '取消'),
      )
    },
  },
])
/**合计行 */
const summaryMethod: DataTableCreateSummary = (pageData) => {
  const totalSize =
    pageData.reduce((prev, row) => prev + (row.size as number), 0) / 1024 / 1024 / 1024
  return {
    name: {
      value: h('span', null, [
        h('span', { class: 'text-gray' }, `模型数量`),
        h('span', { class: 'px-4px' }, pageData.length),
        h('span', { class: 'text-gray' }, '个'),
      ]),
    },
    size: {
      value: h('span', null, [
        h('span', { class: 'text-gray' }, `总大小`),
        h('span', { class: 'px-4px' }, totalSize.toFixed(3)),
        h('span', { class: 'text-gray' }, `GB`),
      ]),
    },
  }
}
/**模型列表 */
const data = ref<ModelList[]>([])
/**获取 */
const getModelList = async () => {
  try {
    const { models = [] } = await ollama.list()
    data.value = models.map((item) => ({
      ...item,
      status: 'success',
      completed: item.size,
      msg: '已下载',
      time: 0,
      speed: 0,
    }))
  } catch (error) {
    message.error('获取模型列表失败，请检查ollama服务是否已启动')
  }
}

/**初始化 */
onMounted(() => {
  getModelList()
})
/**拉取的模型名称 */
const modelName = ref('')
/**正在拉取的模型 */
const pullingList: { name: string; aborter: AbortableAsyncIterator<ProgressResponse> }[] = []

const handlePull = () => {
  if (!modelName.value) {
    message.error('请输入模型名称')
    return
  }
  if (!data.value.some((item) => item.name === modelName.value)) {
    data.value.push({
      name: modelName.value,
      size: 0,
      modified_at: new Date(),
      digest: '',
      expires_at: new Date(),
      size_vram: 0,
      details: {
        parent_model: '',
        format: '',
        family: '',
        families: [],
        parameter_size: '',
        quantization_level: '',
      },
      status: 'success',
      msg: '已下载',
      completed: 0,
      time: 0,
      speed: 0,
    })
  }
  pullFromOllama(modelName.value)
  modelName.value = ''
}
const pullFromOllama = async (name: string) => {
  const target = data.value.find((item) => item.name === name)!
  const res = await ollama.pull({ model: name, stream: true })
  target.aborter = res
  try {
    for await (const element of res) {
      const { status, total = 0, completed = 0 } = element
      target.size = total
      if (status === 'success' || target.status !== status) {
        target.status = status
      }
      const time = Date.now()
      if (time - target.time > 1000) {
        target.speed = completed - target.completed
        target.completed = completed
        target.time = time
      }
    }
  } catch (error) {
    const { message } = error as Error
    if (message.includes('aborted')) {
      target.status = 'cancel'
    } else {
      target.status = 'error'
      target.msg = message
    }
  }
  getModelList()
}
</script>
<style lang="scss"></style>
